
from datetime import datetime

from agentscope.agent import ReActAgent, UserAgent
from agentscope.model import DashScopeChatModel, OpenAIChatModel
from agentscope.formatter import DashScopeChatFormatter, OpenAIChatFormatter
from agentscope.memory import InMemoryMemory
from agentscope.tool import Toolkit, execute_python_code, execute_shell_command
import os, asyncio
from utils import model
import agentscope


current_date = datetime.now().strftime("%Y年%m月%d日")
sys_prompts = f'''
AI智能量化投资助手。
擅长A股，港股，美股，ETF，期货等量化投资。
当前日期：{current_date}
'''

from tools.mcp_utils import map_mcp

async def main():
    toolkit = Toolkit()
    toolkit.register_tool_function(execute_python_code)
    toolkit.register_tool_function(search)
    toolkit.register_tool_function(execute_shell_command)

    # 从 MCP 服务器注册所有工具
    await toolkit.register_mcp_client(
        map_mcp,
        # group_name="map_services",  # 可选的组名
    )

    print("注册的 MCP 工具总数：", len(toolkit.get_json_schemas()))


    agent = ReActAgent(
        name="AI量化小助手",
        sys_prompt=sys_prompts,
        model=model,
        memory=InMemoryMemory(),
        formatter=OpenAIChatFormatter(),
        toolkit=toolkit,
    )

    user = UserAgent(name="user")

    msg = None
    while True:
        msg = await agent(msg)
        msg = await user(msg)
        if msg.get_text_content() == "exit":
            break
agentscope.init(studio_url="http://localhost:3000")
asyncio.run(main())